module fft_wrap

!----------------------------------------------
! VARIABLES
!----------------------------------------------

  use, intrinsic :: iso_c_binding

  use dimensions, only : jx, jy, jz, jv, jdir

  include 'fftw3.f03'

 
  integer(kind=8) ::  plan_xforward,  plan_xbackward
  integer(kind=8) ::  plan_yforward,  plan_ybackward
  integer(kind=8) ::  plan_zforward,  plan_zbackward
  integer(kind=8) ::  plan_vforward,  plan_vbackward

  integer(kind=8) ::  plan_xyzforward,  plan_xyzbackward

  type(C_PTR)     ::  plan_v5forward,  plan_v5backward
  type(C_PTR)     ::  plan_z5forward,  plan_z5backward
  type(C_PTR)     ::  plan_y5forward,  plan_y5backward
  type(C_PTR)     ::  plan_x5forward,  plan_x5backward

  type(C_PTR)     ::  plan_z4forward,  plan_z4backward
  type(C_PTR)     ::  plan_y4forward,  plan_y4backward
  type(C_PTR)     ::  plan_x4forward,  plan_x4backward


!----------------------------------------------
! SUBROUTINES
!----------------------------------------------


contains

!-- fftw initialization --

subroutine fft_init(dfu)

  use mpi_wrap

  implicit none

  complex, intent(inout)   :: dfu(jv, jz, jy, jx, jDir)
  complex                  :: fx(jx), fy(jy), fz(jz), fv(jv)
  
  integer(kind=4)          :: rank
  integer(kind=4)          :: hm_rank
  type(FFTW_IODIM)         :: dims(1)
  type(FFTW_IODIM)         :: hm_dims(4)

  integer                  :: iret, tid, nthreads
  integer                  :: omp_get_thread_num,  omp_get_max_threads

 !-----------------------------------------------

  write(msg,*) char(10), ' Initializing fft_wrap module'
  call mpi_wrap_msg

  !$OMP PARALLEL PRIVATE(tid)

  tid = omp_get_thread_num()
  nthreads = omp_get_max_threads()
  !print *, 'Hello World from thread = ', tid, ' out of ', nthreads
      
  !$OMP END PARALLEL

  call dfftw_init_threads(iret)

  call dfftw_plan_dft_1d(plan_xforward, jx, fx, fx, FFTW_FORWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_yforward, jy, fy, fy, FFTW_FORWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_zforward, jz, fz, fz, FFTW_FORWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_vforward, jv, fv, fv, FFTW_FORWARD, FFTW_ESTIMATE)

  call dfftw_plan_dft_1d(plan_xbackward, jx, fx, fx, FFTW_BACKWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_ybackward, jy, fy, fy, FFTW_BACKWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_zbackward, jz, fz, fz, FFTW_BACKWARD, FFTW_ESTIMATE)
  call dfftw_plan_dft_1d(plan_vbackward, jv, fv, fv, FFTW_BACKWARD, FFTW_ESTIMATE)

  call dfftw_plan_with_nthreads(nthreads)

  call dfftw_plan_dft_3d(plan_xyzforward,  jz,jy,jx, dfu,dfu, FFTW_FORWARD,  FFTW_ESTIMATE)
  call dfftw_plan_dft_3d(plan_xyzbackward, jz,jy,jx, dfu,dfu, FFTW_BACKWARD, FFTW_ESTIMATE)



  !-- 1D transform of 5D arrays --

  rank = 1
  hm_rank = 4


  !--  1D transforms of 5D array in v-direction --

     dims(1)%n  = jv
     dims(1)%is = 1
     dims(1)%os = 1

  hm_dims(4)%n  = jz
  hm_dims(4)%is = jv
  hm_dims(4)%os = jv

  hm_dims(3)%n  = jy
  hm_dims(3)%is = jv*jz
  hm_dims(3)%os = jv*jz

  hm_dims(2)%n  = jx
  hm_dims(2)%is = jv*jz*jy
  hm_dims(2)%os = jv*jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jv*jz*jy*jx
  hm_dims(1)%os = jv*jz*jy*jx

  plan_v5forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_v5backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);

 
  !--  1D transforms of 5D array in z-direction --

  hm_dims(4)%n  = jv
  hm_dims(4)%is = 1
  hm_dims(4)%os = 1

     dims(1)%n  = jz
     dims(1)%is = jv
     dims(1)%os = jv

  hm_dims(3)%n  = jy
  hm_dims(3)%is = jv*jz
  hm_dims(3)%os = jv*jz

  hm_dims(2)%n  = jx
  hm_dims(2)%is = jv*jz*jy
  hm_dims(2)%os = jv*jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jv*jz*jy*jx
  hm_dims(1)%os = jv*jz*jy*jx

  plan_z5forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_z5backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);


  !--  1D transforms of 5D array in y-direction --

  hm_dims(4)%n  = jv
  hm_dims(4)%is = 1
  hm_dims(4)%os = 1

  hm_dims(3)%n  = jz
  hm_dims(3)%is = jv
  hm_dims(3)%os = jv

     dims(1)%n  = jy
     dims(1)%is = jv*jz
     dims(1)%os = jv*jz

  hm_dims(2)%n  = jx
  hm_dims(2)%is = jv*jz*jy
  hm_dims(2)%os = jv*jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jv*jz*jy*jx
  hm_dims(1)%os = jv*jz*jy*jx

  plan_y5forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_y5backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);


  !--  1D transforms of 5D array in x-direction --

  hm_dims(4)%n  = jv
  hm_dims(4)%is = 1
  hm_dims(4)%os = 1

  hm_dims(3)%n  = jz
  hm_dims(3)%is = jv
  hm_dims(3)%os = jv

  hm_dims(2)%n  = jy
  hm_dims(2)%is = jv*jz
  hm_dims(2)%os = jv*jz

     dims(1)%n =  jx
     dims(1)%is = jv*jz*jy
     dims(1)%os = jv*jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jv*jz*jy*jx
  hm_dims(1)%os = jv*jz*jy*jx

  plan_x5forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_x5backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);



  !-- 1D transform of 4D arrays --

  rank = 1
  hm_rank = 3

 
  !--  1D transforms of 4D array in z-direction --

     dims(1)%n  = jz
     dims(1)%is = 1
     dims(1)%os = 1

  hm_dims(3)%n  = jy
  hm_dims(3)%is = jz
  hm_dims(3)%os = jz

  hm_dims(2)%n  = jx
  hm_dims(2)%is = jz*jy
  hm_dims(2)%os = jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jz*jy*jx
  hm_dims(1)%os = jz*jy*jx

  plan_z4forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_z4backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);


  !--  1D transforms of 4D array in y-direction --

  hm_dims(3)%n  = jz
  hm_dims(3)%is = 1
  hm_dims(3)%os = 1

     dims(1)%n  = jy
     dims(1)%is = jz
     dims(1)%os = jz

  hm_dims(2)%n  = jx
  hm_dims(2)%is = jz*jy
  hm_dims(2)%os = jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jz*jy*jx
  hm_dims(1)%os = jz*jy*jx

  plan_y4forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_y4backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);


  !--  1D transforms of 4D array in x-direction --

  hm_dims(3)%n  = jz
  hm_dims(3)%is = 1
  hm_dims(3)%os = 1

  hm_dims(2)%n  = jy
  hm_dims(2)%is = jz
  hm_dims(2)%os = jz

     dims(1)%n =  jx
     dims(1)%is = jz*jy
     dims(1)%os = jz*jy

  hm_dims(1)%n  = jdir
  hm_dims(1)%is = jz*jy*jx
  hm_dims(1)%os = jz*jy*jx

  plan_x4forward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_FORWARD, FFTW_ESTIMATE);

  plan_x4backward = fftw_plan_guru_dft(rank, dims, hm_rank, hm_dims,  &
                   dfu, dfu, FFTW_BACKWARD, FFTW_ESTIMATE);


end subroutine fft_init

!----------------------------------------------

end module
